package bjut.core.netty;

import bjut.core.extension.ExtensionLoader;
import bjut.core.factory.SingletonFactory;
import bjut.core.Message.RpcMessage;
import bjut.core.Message.RpcRequest;
import bjut.core.Message.RpcResponse;
import bjut.core.channel.ChannelRecord;
import bjut.core.constant.RpcConstant;
import bjut.core.enums.CompressTypeEnum;
import bjut.core.enums.SerializationTypeEnum;
import bjut.core.interfaces.ServiceDiscovery;
import bjut.core.interfaces.ZRpcRequestTransport;
import bjut.core.processor.UnprocessedRequest;
import io.netty.bootstrap.Bootstrap;
import io.netty.channel.*;
import io.netty.channel.nio.NioEventLoopGroup;
import io.netty.channel.socket.nio.NioSocketChannel;
import io.netty.handler.logging.LogLevel;
import io.netty.handler.logging.LoggingHandler;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;

import java.net.InetSocketAddress;
import java.util.concurrent.CompletableFuture;

/**
 * @Author: zhuzhiming
 * @Date: 2022/4/22 14:06
 */
@Slf4j
public class ZRpcNettyClient implements ZRpcRequestTransport {

    private final Bootstrap bootstrap;
    private final EventLoopGroup workerGroup;
    private final ChannelRecord channelRecord;

    private final ServiceDiscovery serviceDiscovery;

    private final UnprocessedRequest unprocessedRequest;


    public ZRpcNettyClient() {
        workerGroup = new NioEventLoopGroup();
        bootstrap = new Bootstrap();
        bootstrap.group(workerGroup)
                .channel(NioSocketChannel.class)
                .handler(new LoggingHandler(LogLevel.INFO))
                .option(ChannelOption.CONNECT_TIMEOUT_MILLIS, 5000)
                .handler(new NettyClientChannelInitializer());

        this.channelRecord = SingletonFactory.getInstance(ChannelRecord.class);
        this.serviceDiscovery = ExtensionLoader.getExtensionLoader(ServiceDiscovery.class).getExtension("zookeeper");
        this.unprocessedRequest = SingletonFactory.getInstance(UnprocessedRequest.class);
    }


    @Override
    public Object sendRpcRequest(RpcRequest rpcRequest) {

        CompletableFuture<RpcResponse<Object>> resultFuture = new CompletableFuture<>();
        InetSocketAddress inetSocketAddress = serviceDiscovery.lookupService(rpcRequest);
        Channel channel = getChannel(inetSocketAddress);

        if (channel.isActive()) {
            unprocessedRequest.putUnprocessedRequest(rpcRequest.getRequestId(), resultFuture);
            RpcMessage rpcMessage = RpcMessage.builder()
                    .compress(CompressTypeEnum.GZIP.getCode())
                    .codec(SerializationTypeEnum.HESSIAN.getCode())
                    .messageType(RpcConstant.REQUEST_TYPE)
                    .data(rpcRequest)
                    .build();

            channel.writeAndFlush(rpcMessage).addListener((ChannelFutureListener) future -> {

                if (future.isSuccess()) {
                    log.info("client send message [{}] success", rpcMessage);
                } else {
                    future.channel().close();
                    resultFuture.completeExceptionally(future.cause());
                    log.error("client send message [{}]  fail", rpcMessage);
                }
            });
        } else {
            throw new IllegalStateException();
        }

        return resultFuture;

    }


    public Channel getChannel(InetSocketAddress inetSocketAddress) {
        Channel channel = channelRecord.getChannel(inetSocketAddress);
        if (channel == null) {
            channel = doConnect(inetSocketAddress);
            channelRecord.putChannel(inetSocketAddress, channel);
        }
        return channel;
    }


    @SneakyThrows
    public Channel doConnect(InetSocketAddress inetSocketAddress) {
        CompletableFuture<Channel> completableFuture = new CompletableFuture<>();
        bootstrap.connect(inetSocketAddress).addListener((ChannelFutureListener) future -> {
            if (future.isSuccess()) {
                completableFuture.complete(future.channel());
            } else {
                throw new IllegalStateException();
            }
        });
        return completableFuture.get();
    }

}
